package com.example.pay.util;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.FileInputStream;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.security.KeyFactory;
import java.security.PublicKey;
import java.security.Signature;
import java.security.spec.MGF1ParameterSpec;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.security.spec.PSSParameterSpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.Map;

/**
 * 微信 App 支付调起参数签名验签工具。
 * 使用商户证书的公钥（apiclient_cert.pem）对 RSA(SHA256) 的 paySign 进行验签。
 */
public class WeChatAppPaySignVerifier {

    /**
     * 构造 App 支付签名字符串（使用小写键名，固定顺序）。
     */
    public static String buildSignString(String appId, String partnerId, String prepayId, String nonceStr, String timeStamp) {
        StringBuilder sb = new StringBuilder();
        sb.append("appid=").append(appId)
          .append("&noncestr=").append(nonceStr)
          .append("&package=Sign=WXPay")
          .append("&partnerid=").append(partnerId)
          .append("&prepayid=").append(prepayId)
          .append("&timestamp=").append(timeStamp);
        return sb.toString();
    }

    /**
     * v3 App 调起支付的 RSA 签名串（换行拼接，顺序为 appId, timeStamp, nonceStr, prepayId，末尾包含一个换行）。
     */
    public static String buildSignStringAppV3(String appId, String nonceStr, String timeStamp, String prepayId) {
        StringBuilder sb = new StringBuilder();
        sb.append(appId).append('\n')
          .append(timeStamp).append('\n')
          .append(nonceStr).append('\n')
          .append(prepayId).append('\n');
        return sb.toString();
    }

    /**
     * 对 paySign 进行验签。
     */
    public static boolean verify(String signString, String paySignBase64, PublicKey merchantPublicKey) {
        try {
            Signature sig = Signature.getInstance("SHA256withRSA");
            sig.initVerify(merchantPublicKey);
            sig.update(signString.getBytes(StandardCharsets.UTF_8));
            byte[] signBytes = Base64.getDecoder().decode(paySignBase64);
            return sig.verify(signBytes);
        } catch (Exception e) {
            return false;
        }
    }

    /**
     * 使用 RSASSA-PSS（SHA-256, MGF1, saltLen=32）尝试验签，用于诊断是否服务端误用 PSS。
     */
    public static boolean verifyWithRSAPSS(String signString, String paySignBase64, PublicKey merchantPublicKey) {
        try {
            Signature sig = Signature.getInstance("RSASSA-PSS");
            PSSParameterSpec pssSpec = new PSSParameterSpec("SHA-256", "MGF1", MGF1ParameterSpec.SHA256, 32, 1);
            sig.setParameter(pssSpec);
            sig.initVerify(merchantPublicKey);
            sig.update(signString.getBytes(StandardCharsets.UTF_8));
            byte[] signBytes = Base64.getDecoder().decode(paySignBase64);
            return sig.verify(signBytes);
        } catch (Exception e) {
            return false;
        }
    }

    /**
     * 从 PEM 内容加载公钥：支持 X.509 证书（BEGIN CERTIFICATE）或纯公钥（BEGIN PUBLIC KEY）。
     * 支持 file 路径与 classpath: 前缀。
     */
    public static PublicKey loadPublicKeyFromPem(String location) throws Exception {
        InputStream in;
        if (location != null && location.startsWith("classpath:")) {
            String path = location.substring("classpath:".length());
            in = Thread.currentThread().getContextClassLoader().getResourceAsStream(path);
            if (in == null) throw new IllegalArgumentException("证书/公钥未在 classpath 找到: " + path);
        } else {
            in = new FileInputStream(location);
        }
        String pem;
        try (InputStream is = in) {
            pem = readAllUtf8(is).trim();
        }

        if (pem.contains("BEGIN CERTIFICATE")) {
            CertificateFactory factory = CertificateFactory.getInstance("X.509");
            try (ByteArrayInputStream bais = new ByteArrayInputStream(pem.getBytes(StandardCharsets.UTF_8))) {
                X509Certificate cert = (X509Certificate) factory.generateCertificate(bais);
                return cert.getPublicKey();
            }
        }

        if (pem.contains("BEGIN PUBLIC KEY") || pem.contains("BEGIN RSA PUBLIC KEY")) {
            String base64 = pem.replaceAll("-----BEGIN (.*?)-----", "")
                               .replaceAll("-----END (.*?)-----", "")
                               .replaceAll("\\s", "");
            byte[] keyBytes = Base64.getDecoder().decode(base64);
            KeyFactory kf = KeyFactory.getInstance("RSA");
            X509EncodedKeySpec spec = new X509EncodedKeySpec(keyBytes);
            return kf.generatePublic(spec);
        }

        throw new IllegalArgumentException("不支持的 PEM 内容，需为证书或公钥格式");
    }

    private static String readAllUtf8(InputStream is) throws Exception {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        byte[] buf = new byte[4096];
        int n;
        while ((n = is.read(buf)) != -1) {
            baos.write(buf, 0, n);
        }
        return new String(baos.toByteArray(), StandardCharsets.UTF_8);
    }

    /**
     * 直接用 JSON 对象进行验签（payload 节点）。
     */
    public static boolean verifyFromJsonPayload(JsonNode payload, PublicKey merchantPublicKey) {
        String appId = payload.path("appId").asText();
        String partnerId = payload.path("partnerId").asText();
        String prepayId = payload.path("prepayId").asText();
        String nonceStr = payload.path("nonceStr").asText();
        String timeStamp = payload.path("timeStamp").asText();
        String paySign = payload.path("paySign").asText();
        // 先尝试 v3 换行签名串
        String v3 = buildSignStringAppV3(appId, nonceStr, timeStamp, prepayId);
        if (verify(v3, paySign, merchantPublicKey)) return true;
        // 再回退到老的键值拼接（部分历史实现可能仍使用此方式）
        String kv = buildSignString(appId, partnerId, prepayId, nonceStr, timeStamp);
        return verify(kv, paySign, merchantPublicKey);
    }

    /**
     * 尝试多种常见 App 支付签名串排列进行验签（容错）。
     */
    public static boolean verifyFromJsonPayloadVariants(JsonNode payload, PublicKey merchantPublicKey) {
        String appId = payload.path("appId").asText();
        String partnerId = payload.path("partnerId").asText();
        String prepayId = payload.path("prepayId").asText();
        String nonceStr = payload.path("nonceStr").asText();
        String timeStamp = payload.path("timeStamp").asText();
        String paySign = payload.path("paySign").asText();

        // 先尝试 v3 换行签名串
        String v3line = buildSignStringAppV3(appId, nonceStr, timeStamp, prepayId);
        if (verify(v3line, paySign, merchantPublicKey)) return true;

        String v1 = buildSignString(appId, partnerId, prepayId, nonceStr, timeStamp);
        if (verify(v1, paySign, merchantPublicKey)) return true;

        // 变体2：将 partnerid、prepayid 提前
        String v2 = new StringBuilder()
                .append("appid=").append(appId)
                .append("&partnerid=").append(partnerId)
                .append("&prepayid=").append(prepayId)
                .append("&package=Sign=WXPay")
                .append("&noncestr=").append(nonceStr)
                .append("&timestamp=").append(timeStamp)
                .toString();
        if (verify(v2, paySign, merchantPublicKey)) return true;

        // 变体3：按键名字典序（appid, noncestr, package, partnerid, prepayid, timestamp）与 v1 相同，这里仍保留尝试
        String[] keys = {"appid","noncestr","package","partnerid","prepayid","timestamp"};
        String[] values = {appId, nonceStr, "Sign=WXPay", partnerId, prepayId, timeStamp};
        StringBuilder v3sb = new StringBuilder();
        for (int i = 0; i < keys.length; i++) {
            if (i > 0) v3sb.append('&');
            v3sb.append(keys[i]).append('=').append(values[i]);
        }
        String v3 = v3sb.toString();
        return verify(v3, paySign, merchantPublicKey);
    }

    /**
     * 解析 JSON 字符串并进行验签（包含 channel/orderCode/payload 结构）。
     */
    public static boolean verifyFromJson(String json, PublicKey merchantPublicKey) throws Exception {
        ObjectMapper mapper = new ObjectMapper();
        JsonNode root = mapper.readTree(json);
        JsonNode payload = root.path("payload");
        return verifyFromJsonPayloadVariants(payload, merchantPublicKey);
    }
}